Skip to content

Migrate to stablehlo #1921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 77 commits into from
Aug 20, 2025
Merged

Migrate to stablehlo #1921

merged 77 commits into from
Aug 20, 2025

Conversation

@paul0403
Copy link
Member

@paul0403 paul0403 commented Jul 18, 2025 •
edited
Loading

Context:
Migrate to stablehlo.

Catalyst depends on the stablehlo dialect because that is what jax lowers classical primitives (additions, jnp array ops, linalg ops, etc) to. The stablehlo dialect was pulled in from the tensorflow/mlir-hlo repo, but that repo is slated for deletion .
They suggest users to migrate to the stablehlo repo.

Description of the Change:
Use the stablehlo version of the mhlo operations, dialects and passes we are currently using.
Remove mhlo build system from regular makefiles and CI scripts.
Remove mhlo patches.

Benefits:
No longer depending on a sunsetting external project.

[sc-96641]

@paul0403
Copy link
Member Author

paul0403 commented Jul 21, 2025 •
edited
Loading

I'm opening CI to check the build system. As of this commit (ee515f5), everything should succeed up until the python tests (i.e. everything should succeed up until make test-mlir).

I'm still building stablehlo in embedded mode instead of external mode, but the strange logic in stablehlo's cmakelists makes it seem impossible to use their external mode. Embedded mode, aka building the stablehlo directory just as another subdirectory in the overall mlir directory, is just the easiest mode to get things going for now.

I will now move on to check the missing passes. We can revisit build mode later.

@paul0403 paul0403 marked this pull request as ready for review July 21, 2025 14:46
@dime10
Copy link
Contributor

dime10 commented Jul 21, 2025

I'm opening CI to check the build system. As of this commit (ee515f5), everything should succeed up until the python tests (i.e. everything should succeed up until make test-mlir).

I'm still building stablehlo in embedded mode instead of external mode, but the strange logic in stablehlo's cmakelists makes it seem impossible to use their external mode. Embedded mode, aka building the stablehlo directory just as another subdirectory in the overall mlir directory, is just the easiest mode to get things going for now.

I will now move on to check the missing passes. We can revisit build mode later.

Thanks @paul0403! What happened to the stdlib error?

@paul0403
Copy link
Member Author

I'm opening CI to check the build system. As of this commit (ee515f5), everything should succeed up until the python tests (i.e. everything should succeed up until make test-mlir).
I'm still building stablehlo in embedded mode instead of external mode, but the strange logic in stablehlo's cmakelists makes it seem impossible to use their external mode. Embedded mode, aka building the stablehlo directory just as another subdirectory in the overall mlir directory, is just the easiest mode to get things going for now.
I will now move on to check the missing passes. We can revisit build mode later.

Thanks @paul0403! What happened to the stdlib error?

The error message is from cpp20. I downgraded to cpp17 and it went away. This is not ideal but at least this frees me up to investigate the passes for now.

paul0403 added 3 commits July 21, 2025 16:51
@paul0403
add back stablehlo legalize to std pass
ae895d0
@paul0403
add back legalize-to-control-flow pass
c938853
@paul0403
register stablehlo optimization passes
c18f678 
This is needed for `stablehlo-aggressive-simplification`, which contains the lowering for stablehlo::ReduceOp
@paul0403
Copy link
Member Author

Needs to wait on llvm update (#1916). stablehlo.concatenate lowers to tensor.concat which has its bufferization patterns implemented after the current llvm version we track.

@paul0403 paul0403 requested a review from dime10 August 14, 2025 19:46
dime10
dime10 reviewed Aug 14, 2025
View reviewed changes
dime10
dime10 approved these changes Aug 14, 2025
View reviewed changes
Copy link
Contributor

@dime10 dime10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't believe it finally happened 😍

dime10
dime10 reviewed Aug 14, 2025
View reviewed changes
@paul0403 paul0403 mentioned this pull request Aug 18, 2025
Merged
@rturrado
Copy link
Contributor

rturrado commented Aug 19, 2025 •
edited
Loading

Regarding the current build wheels errors:

  1. Could it be that --workflow is missing from the list of CMAKE_ARGS (catalyst/runtime/CMakeLists.txt, ExternalProject_Add(lapacke-accelerate? See current Accelerate LAPACKE README.
  2. Would it be worth considering sticking to a given working GIT_TAG (version, UUID...), then doing regular updates for these external projects?

@paul0403
Copy link
Member Author

paul0403 commented Aug 19, 2025 •
edited
Loading

Regarding the current build wheels errors:

  1. Could it be that --workflow is missing from the list of CMAKE_ARGS (catalyst/runtime/CMakeLists.txt, ExternalProject_Add(lapacke-accelerate? See current Accelerate LAPACKE README.
  2. Would it be worth considering sticking to a given working GIT_TAG (version, UUID...), then doing regular updates for these external projects?

The current failures are caused by github mac runner images being upgraded to macos 15. We are handling it in #1986

@paul0403 paul0403 merged commit 449df62 into main Aug 20, 2025
58 of 59 checks passed
@paul0403 paul0403 deleted the migrate_to_stablehlo branch August 20, 2025 13:46
@joeycarter joeycarter mentioned this pull request Oct 22, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@dime10 dime10 dime10 approved these changes

@mehrdad2m mehrdad2m mehrdad2m left review comments

@mlxd mlxd Awaiting requested review from mlxd

Assignees

No one assigned

Labels

author:build-wheels Run the wheel building workflows on this Pull Request reviewer:require-wheels Pull Requests will need wheel building job successful before being merged

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@paul0403 @dime10 @rturrado @mehrdad2m